"""
主要定义了和区块链的交互
"""
import web3.exceptions
from web3 import Web3
from fl.config import Config
from fl.util.ipfs import IPFSManager
from web3.middleware import geth_poa_middleware
from chain.contracts import model_update_registry as mur
from fl.logger import logger


# 区块链管理器
class BlockchainManager:
    def __init__(self, config: Config, ipfs_client: IPFSManager=None):

        self.config = config
        self.ipfs_client = ipfs_client
        self.w3 = Web3(Web3.HTTPProvider(config.ETHEREUM_HTTP_PROVIDER))
        self.w3.middleware_onion.inject(geth_poa_middleware, layer=0)
        self.contract = self.w3.eth.contract(address=mur.CONTRACT_ADDRESS, abi=mur.ABI)
        self.account = self.w3.eth.accounts[0]

    # 提交局部模型
    def submit_local_model_update(self, ipfs_hash, num_examples, train_round):
        tx = self.contract.functions.submitLocalUpdate(ipfs_hash, num_examples, train_round).transact({
            'from': self.account,
            'nonce': self.w3.eth.get_transaction_count(self.account)
        })
        receipt = self.w3.eth.wait_for_transaction_receipt(tx)
        if receipt['status'] == 1:
            logger.info("Submit local models update successfully!")
        else:
            logger.error("Submit local models update failed!")

    # 提交全局模型
    def submit_global_model_update(self, ipfs_hash, min_member):
        try:
            tx = self.contract.functions.submitGlobalUpdate(ipfs_hash, min_member).transact({
                'from': self.account,
                'nonce': self.w3.eth.get_transaction_count(self.account)
            })
            receipt = self.w3.eth.wait_for_transaction_receipt(tx)
            if receipt['status'] == 1:
                logger.info("Submit global models update successfully!")
            else:
                logger.error("Submit global models update failed!")
        except Exception as e:
            logger.error(e.args[0])

    def submit_candidate_global_model_update(self, ipfs):
        try:
            tx = self.contract.functions.submitCandidateGlobalUpdate(ipfs).transact({
                'from': self.account,
                'nonce': self.w3.eth.get_transaction_count(self.account)
            })
            receipt = self.w3.eth.wait_for_transaction_receipt(tx)
            if receipt['status'] == 1:
                logger.info("Submit candidate global models update successfully!")
            else:
                logger.error("Submit candidate global models update failed!")
        except Exception as e:
            logger.error(e.args[0])

    def get_candidate_global_model_update(self):
        try:
            candidate_global_model_updates = self.contract.functions.getCandidateGlobalUpdates().call()
            return candidate_global_model_updates
        except web3.exceptions.ContractLogicError as e:
            logger.warning(e.args[0])

    def get_candidate_aggregators(self):
        try:
            candidate_aggregators = self.contract.functions.getSelectedAggregators().call()
            return candidate_aggregators
        except web3.exceptions.ContractLogicError as e:
            logger.warning(e.args[0])

    def vote_global_model_update(self, uploader_addr):
        try:
            tx = self.contract.functions.voteGlobalUpdate(uploader_addr).transact({
                'from': self.account,
                'nonce': self.w3.eth.get_transaction_count(self.account)
            })
            receipt = self.w3.eth.wait_for_transaction_receipt(tx)
            if receipt['status'] == 1:
                logger.info("Vote global models update successfully!")
            else:
                logger.error("Vote global models update failed!")

        except Exception as e:
            logger.error(e.args[0])

    # 选举聚合节点
    def select_aggregators(self):
        try:
            tx = self.contract.functions.selectAggregators().transact({
                'from': self.account,
                'nonce': self.w3.eth.get_transaction_count(self.account)
            })
            self.w3.eth.wait_for_transaction_receipt(tx)
        except web3.exceptions.ContractLogicError as e:
            logger.warning(e.args[0])

    # 初始化模型，当区块链中不存在初始模型时执行
    def initialize_model(self, model):
        initial_parameters = [val.cpu().numpy() for val in model.state_dict().values()]
        initial_ipfs_hash = self.ipfs_client.upload_to_ipfs(initial_parameters)
        logger.info(f"Initializing model{initial_ipfs_hash}:{self.config.MIN_MEMBERS}")
        self.submit_global_model_update(initial_ipfs_hash, self.config.MIN_MEMBERS)
        logger.info(f"Submitted initial global models to IPFS hash: {initial_ipfs_hash}")

    # 设置最大聚合轮数
    def set_max_round(self, max_round_of_fl):
        try:
            tx = self.contract.functions.setMaxRound(max_round_of_fl).transact({
                'from': self.account,
                'nonce': self.w3.eth.get_transaction_count(self.account)
            })
            receipt = self.w3.eth.wait_for_transaction_receipt(tx)
            if receipt['status'] == 1:
                logger.info("Set max round of FL successfully!")
            else:
                logger.error()
        except web3.exceptions.ContractLogicError as e:
            logger.warning(e.args[0])

    def refresh_events(self):
        tx = self.contract.functions.refreshEvents().transact({
            'from': self.account,
            'nonce': self.w3.eth.get_transaction_count(self.account)
        })
        self.w3.eth.wait_for_transaction_receipt(tx)